package org.example.ycsb;

import org.apache.commons.lang3.RandomStringUtils;
import org.apache.commons.lang3.StringUtils;
import org.apache.flink.api.common.functions.RuntimeContext;
import org.apache.flink.api.java.utils.ParameterTool;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.functions.source.RichParallelSourceFunction;
import org.apache.hadoop.hbase.util.MD5Hash;

import java.security.MessageDigest;
import java.util.ArrayList;
import java.util.Base64;
import java.util.List;

public class MockSource extends RichParallelSourceFunction<ProcessEvent> {
    private int indexOfThisSubTask;
    private int numberOfParallelSubTasks;
    private long upper;
    private long lower;
    private transient volatile boolean flag;
    private ConfigBean configBean;

    @Override
    public void open(Configuration parameters) throws Exception {
        final RuntimeContext runtimeContext = getRuntimeContext();
        this.indexOfThisSubTask = runtimeContext.getIndexOfThisSubtask();
        this.numberOfParallelSubTasks = runtimeContext.getNumberOfParallelSubtasks();
        ParameterTool parameterTool = ParameterTool.fromMap(runtimeContext.getExecutionConfig().getGlobalJobParameters().toMap());
        String ycsbConfig = parameterTool.get("configYcsbBase64");
        if (ycsbConfig == null || ycsbConfig.trim().isEmpty()) {
            configBean = new ConfigBean();
        } else {
            byte[] base64decodedBytes = Base64.getDecoder().decode(ycsbConfig);
            ycsbConfig = new String(base64decodedBytes, "utf-8");
            configBean = JackSonUtil.readValue(ycsbConfig, ConfigBean.class);
        }
        long rowNum = configBean.getRowNum();
        long bucketSize = rowNum / numberOfParallelSubTasks;
        lower = indexOfThisSubTask * bucketSize;
        upper = lower + bucketSize;
        if (indexOfThisSubTask == numberOfParallelSubTasks - 1) {
            upper = rowNum;
        }
        flag=true;
    }

    @Override
    public void run(SourceContext<ProcessEvent> sourceContext) throws Exception {
        for (long i = lower; i <= upper; i++) {
            if(!flag){
                break;
            }
            ProcessEvent processEvent = new ProcessEvent();
            String str = String.valueOf(i);
            // 此 MessageDigest 类为应用程序提供信息摘要算法的功能
            MessageDigest md5 = MessageDigest.getInstance("MD5");
            // 转换为MD5码
            byte[] digest = md5.digest(str.getBytes("utf-8"));
            String hexStr = bytesToHex(digest);
            if(configBean.getRowKeyLength()>32){
                hexStr = StringUtils.rightPad(hexStr,configBean.getRowKeyLength(),"0");
            }else{
                hexStr = hexStr.substring(0,configBean.getRowKeyLength());
            }
            processEvent.setRowKey(hexStr);
            List<Integer> columnSizes=  configBean.getColumnSizes();
            List<String> columnDatas = new ArrayList<>();
            for (Integer columnSize : columnSizes) {
                columnDatas.add(RandomStringUtils.randomAlphanumeric(columnSize));
            }
            processEvent.setValues(columnDatas);
            if(configBean.isInitMode()){
                processEvent.setFlag(1);
            }else{
                double readRatio = configBean.getReadRatio();
                if(Math.random()<=readRatio){
                    processEvent.setFlag(2);
                }else{
                    if(configBean.isReadModifyWrite()){
                        processEvent.setFlag(3);
                    }else{
                        processEvent.setFlag(1);
                    }
                }
            }
            sourceContext.collect(processEvent);
        }
    }

    @Override
    public void cancel() {
        flag=false;
    }

    private  final char[] HEX_ARRAY = "0123456789ABCDEF".toCharArray();

    public  String bytesToHex(byte[] bytes) {
//        char[] hexChars = new char[bytes.length * 2];
//        for (int j = 0; j < bytes.length; j++) {
//            int v = bytes[j] & 0xFF;
//            hexChars[j * 2] = HEX_ARRAY[v >>> 4];
//            hexChars[j * 2 + 1] = HEX_ARRAY[v & 0x0F];
//        }
        return MD5Hash.getMD5AsHex(bytes);
//        return new String(hexChars);
    }


    public String bytesToHexString(byte[] bArray, int length) {
        StringBuffer sb = new StringBuffer(length);
        String sTemp;
        for (int i = 0; i < length; i++) {
            sTemp = Integer.toHexString(0xFF & bArray[i]);
            if (sTemp.length() < 2)
                sb.append(0);
            sb.append(sTemp.toUpperCase());
        }
        return sb.toString();
    }
}
